#include <cuda_runtime.h>
#include <cuda_fp16.h> 
#include "conv2d.h"

extern "C" __global__ void implicit_gemm_v1(mykernelParamType param)
{ 
    // 三维线程映射：X-输出位置，Y-输出通道，Z-批处理
    int OhOw = blockIdx.x * blockDim.x + threadIdx.x;
    int k = blockIdx.y * blockDim.y + threadIdx.y;
    int n = blockIdx.z;
    // 边界检查
    if(OhOw >= param.Oh*param.Ow || k >= param.k || n >= param.n)
        return;

    // 分解输出坐标
    int oh = OhOw / param.Ow;
    int ow = OhOw % param.Ow;
    int input_addr, weight_addr, output_addr;
    float sum = 0.0;
    
    // 输出内存地址（NCHW布局）
    output_addr = n * param.k * param.Oh * param.Ow + k * param.Oh * param.Ow + oh * param.Ow + ow;

    // 合并C/R/S维度为单一循环
    for(int crs = 0; crs < param.c*param.r*param.s; crs++) 
    {
        // 分解三维索引
        int c = crs / (param.r * param.s);
        int residual = crs % (param.r * param.s);
        int r = residual / param.s;
        int s = residual % param.s;
        // 计算输入坐标
        int ih = oh * param.u - param.p + r;
        int iw = ow * param.v - param.q + s;
        if (ih >= 0 && ih < param.h && iw >= 0 && iw < param.w) 
        {
            // 输入内存地址（NCHW布局）
            input_addr = n * param.c * param.h * param.w + c * param.h * param.w + ih * param.w + iw;
            // 权重内存地址（KCRS布局）
            weight_addr = k * param.c * param.r * param.s + c * param.r * param.s + r * param.s + s;
            sum += param.pin[input_addr] * param.pweight[weight_addr];
        }
    }
    param.pout[output_addr] = sum;
}
// 核函数启动配置
void launch_implicit_gemm_v1(unsigned int outh, unsigned int outw, 
                           unsigned int k, unsigned int n, 
                           mykernelParamType* param) {
    // 网格维度计算（向上取整）
    int blockx = (outh * outw + 15) / 16;  // X维度处理输出位置
    int blocky = (k + 15) / 16;            // Y维度处理输出通道
    int blockz = n;                        // Z维度处理批处理
    dim3 block(16, 16, 1);  // 固定线程块尺寸
    dim3 grid(blockx, blocky, blockz);
    implicit_gemm_v1<<<grid, block>>>(*param);
}
